#include <crhuffmanex.h>
using namespace crhuffman;

struct header_mask
{
	static const byte tiny_sizes = 0x20;
	static const byte half_bytes = 0x10;
	static const byte no_huffman = 0x08;
	static const byte bi_of_size = 0x07;
};

#define TINYMODE
#ifdef TINYMODE
#endif

inline void prepare(huffman_builder& builder, istream_base& is)
{
	try {
		while (true) {
			byte ch = is.getbyte();
			builder.push(ch);
		}
	}
	catch (exception<stream_base>& e)
	{
		if (e.code == istream_base::eos && e.obj == &is)
		{}
		else CR_Prompt(_T("builder : istream_base raised unknown exception"));
	}
}

inline void compress(huffman_builder& builder, istream_base& is, ostream_base& os)
{
	huffman_encoder encoder;
	char obi;

#ifdef TINYMODE

	try {
		huffman_tree<16> tree;
		frequencies_wrap<16> freqs;
		freq_merge(freqs, builder);
		code_table_wrap<16> codes;
		tree.build(freqs.freqs_);
		tree.generate(codes.codes_);
		huffman_header<16> header;
		bit_size bsize;
		bool tiny_sizes = header.calcsize(
			codes.codes_, freqs.freqs_, bsize);
		byte flag = (header_mask::bi_of_size & bsize.obi);
		flag |= header_mask::half_bytes;
		if (tiny_sizes)
			flag |= header_mask::tiny_sizes;
		os.putbyte(flag);
		os.putuint(bsize.oBy);
		obi = header.writecodes(os, codes.codes_, tiny_sizes);
		code_merge(encoder, codes);
	}
	catch (exception<huffman_tree<16>>& e)
	{
		switch (e.code)
		{
		case huffman_tree<16>::load_damaged_stream:
			CR_Prompt(_T("huffman_tree : load_damaged_stream"));break;
		case huffman_tree<16>::grow_damaged_stream:
			CR_Prompt(_T("huffman_tree : grow_damaged_stream"));break;
		case huffman_tree<16>::step_damaged_stream:
			CR_Prompt(_T("huffman_tree : step_damaged_stream"));break;
		default: CR_Prompt(_T("huffman_tree : raised unknown exception"));
		}
	}

#else // !TINYMODE

	try {
		huffman_tree<256> tree;
		tree.build(builder.freqs_);
		tree.generate(encoder.codes_);
		huffman_header<256> header;
		bit_size bsize;
		bool tiny_sizes = header.calcsize(
			encoder.codes_, builder.freqs_, bsize);
		byte flag = (header_mask::bi_of_size & bsize.obi);
		if (tiny_sizes)
			flag |= header_mask::tiny_sizes;
		os.putbyte(flag);
		os.putuint(bsize.oBy);
		obi = header.writecodes(os, encoder.codes_, tiny_sizes);
	}
	catch (exception<huffman_tree<256>>& e)
	{
		switch (e.code)
		{
		case huffman_tree<256>::load_damaged_stream:
			CR_Prompt(_T("huffman_tree : load_damaged_stream"));break;
		case huffman_tree<256>::grow_damaged_stream:
			CR_Prompt(_T("huffman_tree : grow_damaged_stream"));break;
		case huffman_tree<256>::step_damaged_stream:
			CR_Prompt(_T("huffman_tree : step_damaged_stream"));break;
		default: CR_Prompt(_T("huffman_tree : raised unknown exception"));
		}
	}

#endif //TINYMODE

	TRACE(_T("****** encoding bit stream: ******\n"));

	obstream bitstrm(os, obi);
	encoder.setbstrm(&bitstrm);

	try {
		while (true) {
			byte ch = is.getbyte();
			encoder.push(ch);
		}
	}
	catch (exception<stream_base>& e)
	{
		if (e.code == istream_base::eos && e.obj == &is)
		{}
		else CR_Prompt(_T("builder : istream_base raised unknown exception"));
	}

	bitstrm.final();

	DEBUGING(bitstrm.dump(std::ofstream("g:\\text_enc.txt")));
	TRACE(_T("\n"));

}

template <typename type_decoder>
inline void decompress(type_decoder& decoder, byte flag, istream_base& is)
{
	const bit_addr endpos(is.tell() + is.getuint(),
		flag & header_mask::bi_of_size);

	TRACE(_T("****** load codes: ******\n"));

	code_table_wrap<256> codes;

#ifdef TINYMODE

	ASSERT(flag & header_mask::half_bytes);
	code_table_wrap<16> codes16;
	huffman_header<16> header;
	char obi = header.readcodes(is, codes16.codes_, 
		!!(flag & header_mask::tiny_sizes));
	code_merge(codes, codes16);

#else //!TINYMODE

	ASSERT(!(flag & header_mask::half_bytes));

	huffman_header<256> header;
	char obi = header.readcodes(is, codes.codes_, 
		!!(flag & header_mask::tiny_sizes));

#endif //TINYMODE

	ibstream bitstrm(is, obi);
	decoder.setbstrm(&bitstrm);

	try {
		TRACE(_T("****** restore tree: ******\n"));
		decoder.restore(codes.codes_);
		TRACE(_T("\n"));
		TRACE(_T("****** decoding bit stream: ******\n"));
		decoder.bitscan(endpos);
		DEBUGING(bitstrm.dump(std::ofstream("g:\\text_dec.txt")));
		TRACE(_T("\n"));
	}
	catch (exception<huffman_tree<256>>& e)
	{
		switch (e.code)
		{
		case huffman_tree<256>::load_damaged_stream:
			CR_Prompt(_T("huffman_tree : load_damaged_stream"));break;
		case huffman_tree<256>::grow_damaged_stream:
			CR_Prompt(_T("huffman_tree : grow_damaged_stream"));break;
		case huffman_tree<256>::step_damaged_stream:
			CR_Prompt(_T("huffman_tree : step_damaged_stream"));break;
		default: CR_Prompt(_T("huffman_tree : raised unknown exception"));
		}
	}

	bitstrm.final(); 

}

int _tmain(int argc, _TCHAR* argv[])
{
	static const char testdoc[] = 
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111"
		"0123456789/*aa00000000000000000000000011111111111111";

	static const size_t doclen = strlen(testdoc)+1;

	byte encbuf[4096];
	byte decbuf[4096];

	printf("****** origin document: ******\n");
	printf(testdoc);
	printf("\n");

	istream_base is;
	ostream_base os;

	huffman_builder builder;
	is.bind((byte*)testdoc, doclen);
	prepare(builder, is);
	if (!builder.analyse(builder.total()))
		return 0;
	is.bind((byte*)testdoc, doclen);
	os.bind(encbuf, 4096);
	os.putuint(is.remain());
	compress(builder, is, os);

	is.bind(encbuf, os.occupy());
	size_t plain = is.getuint();
	const byte flag = is.getbyte();
	if (flag & header_mask::no_huffman) 
		return 0;

	os.bind(decbuf, 4096);
	direct_pusher direct(os);
	huffman_decoder<direct_pusher> decoder;
	decoder.setnext(&direct);
	decompress<huffman_decoder<direct_pusher>>(decoder, flag, is);

	printf("****** target document: ******\n");
	printf((char*)decbuf);
	printf("\n");

	printf("****** follow document: ******\n");
	direct_pusher(os).input((byte*)"=end=", 5);
	os.tell()[20] = 0;
	printf((char*)os.tell()-5);
	printf("\n");

	system("pause");
	return 0;
}

